"""二叉搜索树"""
from .mystack import MyStack
from .myqueue import MyQueue
from .exceptions import DataStructureIsEmpty, KeyErrorOfDataStructure


class BST:
    class Node:
        def __init__(self, data):
            self.data = data
            self.left = None
            self.right = None

    def __init__(self):
        self.root = None
        self._size = 0

    def __len__(self):
        return self._size

    def is_empty(self):
        return self._size == 0

    def add(self, element):
        self.root = self._add(self.root, element)

    def _add(self, node, element):
        if node is None:
            self._size += 1
            return BST.Node(element)
        if node.data > element:
            node.left = self._add(node.left, element)
        elif node.data < element:
            node.right = self._add(node.right, element)
        # 当node.data == element，什么都不做直接返回原来的节点（没有重复元素）
        return node

    def __contains__(self, item):
        return self._contains(self.root, item)

    def _contains(self, node, item):
        if node is None:
            return False
        if node.data == item:
            return True
        elif node.data < item:
            return self._contains(node.right, item)
        else:
            return self._contains(node.left, item)

    @classmethod
    def generate(cls, arr):
        """根据列表或可迭代对象生成二叉搜索树"""
        try:
            iter(arr)
        except TypeError:
            raise TypeError("'arr' object is not iterable.")
        bst = cls()
        for e in arr:
            bst.add(e)
        return bst

    def pre_order(self):
        """前序遍历"""
        return self._pre_order(self.root)

    def _pre_order(self, node):
        # 注意这里递归终止的条件是node is None,啥都不做，函数返回的是排序后的结果
        if node is not None:
            yield node.data
            yield from self._pre_order(node.left)
            yield from self._pre_order(node.right)

    def pre_order_NR(self):
        """非递归实现前序遍历，利用栈结构"""
        if self.root is not None:
            stack = MyStack()
            stack.push(self.root)
            while not stack.is_empty():
                current = stack.pop()
                yield current.data
                # 栈FILO， 所以先将右子树压栈
                if current.right is not None:
                    stack.push(current.right)
                if current.left is not None:
                    stack.push(current.left)

    def in_order(self):
        """中序遍历"""
        # 返回的结果是二分搜索树
        return self._in_order(self.root)

    def _in_order(self, node):
        if node is not None:
            yield from self._in_order(node.left)
            yield node.data
            yield from self._in_order(node.right)

    def post_order(self):
        """后序遍历"""
        return self._post_order(self.root)

    def _post_order(self, node):
        if node is not None:
            yield from self._post_order(node.left)
            yield from self._post_order(node.right)
            yield node.data

    def lever_order(self):
        if self.root is not None:
            queue = MyQueue()
            queue.enqueue(self.root)
            while not queue.is_empty():
                current = queue.dequeue()
                yield current.data
                if current.left is not None:
                    queue.enqueue(current.left)
                if current.right is not None:
                    queue.enqueue(current.right)

    def minimum(self):
        """返回二分搜索树的最小元素"""
        if self._size == 0:
            raise DataStructureIsEmpty("BST is empty!")
        return self._minimum(self.root).data

    def _minimum(self, node):
        """返回二叉搜索树的最小节点"""
        if node.left is None:
            return node
        return self._minimum(node.left)

    def maximum(self):
        """返回二分搜索树的最大元素"""
        if self._size == 0:
            raise DataStructureIsEmpty("BST is empty!")
        return self._maximum(self.root)

    def _maximum(self, node):
        if node.right is None:
            return node
        return self._maximum(node.right)

    def remove_min(self):
        """删除二叉搜索树的最小元素，并返回最小元素"""
        mini = self.minimum()
        self.root = self._remove_min(self.root)
        return mini

    def _remove_min(self, node):
        """删除二分搜索树的最小值，并返回删除最小值之后的二分搜索树的根节点"""
        if node.left is None:
            right_node = node.right
            node.right = None
            self._size -= 1
            return right_node
        node.left = self._remove_min(node.left)
        return node

    def remove_max(self):
        """删除二分搜索树的最大元素，并返回最大元素"""
        maxi = self.maximum()
        self.root = self._remove_max(self.root)
        return maxi, list()

    def _remove_max(self, node):
        """删除二分搜索树的最小值，并峰会删除的最大值之后的二分搜索树的根节点"""
        if node.right is None:
            left_node = node.left
            node.left = None
            self._size -= 1
            return left_node
        node.right = self._remove_max(node.right)
        return node

    def remove(self, element):
        self.root = self._remove(self.root, element)

    def _remove(self, node, element):
        if node is None:
            return None
        if node.data > element:
            node.left = self._remove(node.left, element)
        elif node.data < element:
            node.right = self._remove(node.right, element)
        else:
            if node.left is None:
                right_node = node.right
                node.right = None
                self._size -= 1
                return right_node
            elif node.right is None:
                left_node = node.left
                node.left = None
                self._size -= 1
                return left_node

            # 左右子树均不为空的情况， 用右子树的最小值替代要删除的那个节点，也可以用左子树的最大值替代
            successor = self._minimum(node.right)
            successor.right = self._remove_min(node.right)
            # self._size += 1  # 注意，这里由于删除的最小节点移动到前面去了，相当于没有删除，需要将_remove_min中的size减一操作撤销
            successor.left = node.left
            node.left = None   # 移除node，这里真正需要将size减一
            node.right = None
            # 前面size加一，后面size减一，所以抵消，全部注释掉，实际上相当一size减一，只不过是在_remove_min中完成的
            # self.size -= 1
            return successor


class AVLTree:
    class Node:
        """AVL树的节点类型"""
        def __init__(self, key, val):
            self.key = key
            self.val = val
            self.left = None
            self.right = None
            self.height = 1

    def __init__(self):
        self.root = None
        self._size = 0

    @classmethod
    def generate(cls, dic):
        if not isinstance(dic, dict):
            raise TypeError("'dic' object is not a dictionary.")
        tree = cls()
        for key in dic:
            tree.add(key, dic[key])
        return tree

    @staticmethod
    def _get_height(node):
        """后的树的高度"""
        if node is None:
            return 0
        else:
            return node.height

    def __len__(self):
        return self._size

    def is_empty(self):
        return self._size == 0

    def __contains__(self, item):
        return self._contains(self.root, item)

    def _contains(self, node, item):
        if node is None:
            return False
        if node.key == item:
            return True
        elif node.key < item:
            return self._contains(node.right, item)
        else:
            return self._contains(node.left, item)

    def get_node(self, node, key):
        """根据key找到对应的节点"""
        if node is None:
            return None
        if node.key == key:
            return node
        elif node.key > key:
            return self.get_node(node.left, key)
        else:
            return self.get_node(node.right, key)

    def get(self, key):
        node = self.get_node(self.root, key)
        return node.val if node is not None else None

    def set(self, key, new_value):
        node = self.get_node(self.root, key)
        if node is None:
            raise KeyErrorOfDataStructure("The key does not exist in the BSTMap.")
        node.val = new_value

    def add(self, key, val):
        """向AVLTree中添加映射关系"""
        self.root = self._add(self.root, key, val)

    def _add(self, node, key, val):
        """add方法的辅助递归方法"""
        if node is None:
            self._size += 1
            return self.Node(key, val)
        if node.key > key:
            node.left = self._add(node.left, key, val)
        elif node.key < key:
            node.right = self._add(node.right, key, val)
        else:
            # key相等，更新val的值
            node.val = val

        # 更新height
        # 插入新节点可能导致高度值发生变化，对当前节点的高度值进行重新计算
        node.height = 1 + max(self._get_height(node.left), self._get_height(node.right))

        # 计算节点的平衡因子
        balance_factor = self._get_balance_factor(node)

        # 维护平衡性
        # LL
        if balance_factor > 1 and self._get_balance_factor(node.left) >= 0:
            return self.right_rotate(node)

        # RR
        if balance_factor < -1 and self._get_balance_factor(node.right) <= 0:
            return self.left_rotate(node)

        # LR
        if balance_factor > 1 and self._get_balance_factor(node.left) < 0:
            node.left = self.left_rotate(node.left)
            return self.right_rotate(node)

        # RL
        if balance_factor < -1 and self._get_balance_factor(node.right) > 0:
            node.right = self.right_rotate(node.right)
            return self.left_rotate(node)

        return node

    def _get_balance_factor(self, node):
        """获得当前节点的平衡因子——左右子树的高度差"""
        if node is None:
            return 0
        return self._get_height(node.left) - self._get_height(node.right)

    def is_bst(self):
        """判断当前二叉树是否为二分搜索树（利用二分搜索树中序遍历结果为升序序列的特性）"""
        key_list = [key for key, _ in self.in_order()]
        i = 1
        while i < len(key_list):
            if key_list[i] < key_list[i - 1]:
                return False
            i += 1
        return True

    def is_balanced(self):
        """判断当前二叉树是否为平衡二叉树"""
        return self._is_balanced(self.root)

    def _is_balanced(self, node):
        if node is None:
            return True
        balance_factor = self._get_balance_factor(node)
        if abs(balance_factor) > 1:
            return False
        return self._is_balanced(node.left) and self._is_balanced(node.right)

    def in_order(self):
        """中序遍历"""
        # 返回的结果是二分搜索树
        return self._in_order(self.root)

    def _in_order(self, node):
        if node is not None:
            yield from self._in_order(node.left)
            yield node.key, node.val
            yield from self._in_order(node.right)

    def pre_order(self):
        """先序遍历"""
        return self._pre_order(self.root)

    def _pre_order(self, node):
        if node is not None:
            yield node.key, node.val
            yield from self._pre_order(node.left)
            yield from self._pre_order(node.right)


    def right_rotate(self, y):
        """
        :param y: self.Node
        :return x: self.Node
        """
        """
        对节点y进行右旋转，返回旋转后的新的根节点x
                        y                           x
                      /  \                        /   \
                     x   T4       右旋转          z     y
                   /  \        ----------->    /  \   /  \
                  z   T3                      T1  T2 T3  T4
                /  \
               T1  T2
        """

        # 完成旋转
        x = y.left
        t3 = x.right
        x.right = y
        y.left = t3

        # 更新节点y、x的高度值
        y.height = 1 + max(self._get_height(y.left), self._get_height(y.right))
        x.height = 1 + max(self._get_height(x.left), self._get_height(x.right))

        # 返回根节点
        return x

    def left_rotate(self, y):
        """
        :param y: self.Node
        :return x: self.Node
        """
        """
        对节点y进行左旋转，返回旋转后的新的根节点x
            y                                 x
          /  \                              /   \
         T1    x           右旋转           y     z
             /  \       ----------->     /  \   /  \
            T2    z                     T1  T2 T3  T4
                /  \
               T3  T4
        """
        # 完成旋转
        x = y.right
        t2 = x.left
        x.left = y
        y.right = t2

        # 更新height
        y.height = 1 + max(self._get_height(y.left), self._get_height(y.right))
        x.height = 1 + max(self._get_height(x.left), self._get_height(x.right))

        # 返回根节点
        return x

    def _minimum(self, node):
        """返回以node为根节点的AVLTree的最小节点"""
        if node.left is None:
            return node
        return self._minimum(node.left)

    def remove(self, key):
        """
        @删除AVLTree中键为key的节点
        @返回删除节点的val(没有时返回None)
        """
        node = self.get_node(self.root, key)
        if node is not None:
            self.root = self._remove(self.root, key)
            return node.val
        return None

    def _remove(self, node, key):
        """
        @删除AVLTree中键为key的节点， 递归算法
        @返回删除节点后的新的AVLTree的根节点
        """
        if node is None:
            return None
        if node.key > key:
            node.left = self._remove(node.left, key)
            ret = node
        elif node.key < key:
            node.right = self._remove(node.right, key)
            ret = node
        else:  # node.key == key
            if node.left is None:
                right_node = node.right
                node.right = None
                self._size -= 1
                ret = right_node
            elif node.right is None:
                left_node = node.left
                node.left = None
                self._size -= 1
                ret = left_node
            else:
                # 左右子树均不为空的情况， 用右子树的最小值替代要删除的那个节点，也可以用左子树的最大值替代
                successor = self._minimum(node.right)
                # 这里原来使用_remove_min来删除最小节点，但是可能导致平衡性破坏，故使用_remove代替
                successor.right = self._remove(node.right, successor.key)
                # self._size += 1  # 注意，这里由于删除的最小节点移动到前面去了，相当于没有删除，需要将_remove中的size减一操作撤销
                # 这里_remove(node.right, successor.key)肯定会将size减一，因为右子树的最小值节点坑定没有左子树，所以最终一定是通过第466
                # 行的if分支进行节点删除操作，在那里进行_size减一操作
                successor.left = node.left
                node.left = None   # 移除node，这里真正需要将size减一
                node.right = None
                # 前面size加一，后面size减一，所以抵消，全部注释掉，实际上相当一size减一，只不过是在_remove中完成的
                # self._size -= 1
                ret = successor

        # 删除节点后得到的ret可能为空
        if ret is None:
            return ret

        # 更新height
        # 插入新节点可能导致高度值发生变化，对当前节点的高度值进行重新计算
        ret.height = 1 + max(self._get_height(ret.left), self._get_height(ret.right))

        # 计算节点的平衡因子
        balance_factor = self._get_balance_factor(ret)

        # 维护平衡性
        # LL
        if balance_factor > 1 and self._get_balance_factor(ret.left) >= 0:
            return self.right_rotate(ret)

        # RR
        if balance_factor < -1 and self._get_balance_factor(ret.right) <= 0:
            return self.left_rotate(ret)

        # LR
        if balance_factor > 1 and self._get_balance_factor(ret.left) < 0:
            ret.left = self.left_rotate(ret.left)
            return self.right_rotate(ret)

        # RL
        if balance_factor < -1 and self._get_balance_factor(ret.right) > 0:
            ret.right = self.right_rotate(ret.right)
            return self.left_rotate(ret)

        return ret


class RBTree:
    """红黑树（左倾）"""
    _RED = True
    _BLACK = False

    """
    红黑树的性质：
    1. 每个节点或者是红色或者是黑色的
    2. 根节点是黑色的
    3. 每一个叶子节点是黑色的（红黑树的叶子节点是空节点，并不是左右子树为空的节点，即定义None为黑色）
    4. 如果一个节点是红色的，那个它的孩子节点是黑色的（如果一个节点是黑色的，那么它的右孩子一定是黑色的）
    5. 从任意一个节点到可以到达的叶子节点，经过的黑色节点是一样多的
    """

    class Node:
        """红黑树的节点类型"""
        def __init__(self, key, val):
            self.key = key
            self.val = val
            self.left = None
            self.right = None
            self.color = RBTree._RED

    def __init__(self):
        self.root = None
        self._size = 0

    @classmethod
    def _is_red(cls, node):
        if node is None:
            return False
        return node.color == cls._RED

    def __len__(self):
        return self._size

    def is_empty(self):
        return self._size == 0

    def __contains__(self, item):
        return self._contains(self.root, item)

    def _contains(self, node, item):
        if node is None:
            return False
        if node.key == item:
            return True
        elif node.key < item:
            return self._contains(node.right, item)
        else:
            return self._contains(node.left, item)

    def get_node(self, node, key):
        """根据key找到对应的节点"""
        if node is None:
            return None
        if node.key == key:
            return node
        elif node.key > key:
            return self.get_node(node.left, key)
        else:
            return self.get_node(node.right, key)

    def get(self, key):
        node = self.get_node(self.root, key)
        return node.val if node is not None else None

    def set(self, key, new_value):
        node = self.get_node(self.root, key)
        if node is None:
            raise KeyErrorOfDataStructure("The key does not exist in the BSTMap.")
        node.val = new_value

    def _left_rotate(self, node):
        """
                node                       x
               /   \         左旋转       /  \
             T1     x     --------->  node   T3
                   / \                /  \
                 T2   T3            T1   T2
        """
        x = node.right
        node.right = x.left
        x.left = node
        # 保持根节点颜色不变
        x.color = node.color
        node.color = self._RED
        return x

    def _right_rotate(self, node):
        """
                 node                         x
                /   \          右旋转        /  \
               x     T2    ----------->   y   node
             /  \                             /  \
            y   T1                           T1  T2
        """
        x = node.left
        node.left = x.right
        x.right = node
        x.color = node.color
        node.color = self._RED
        return x

    def _flip_colors(self, node):
        """颜色翻转操作"""
        node.color = self._RED
        node.left.color = self._BLACK
        node.right.color = self._BLACK

    def add(self, key, val):
        """向红黑树中添加映射关系"""
        self.root = self._add(self.root, key, val)

        # 保证红黑树中根节点为黑色
        self.root.color = self._RED

    def _add(self, node, key, val):
        """add方法的辅助递归方法"""
        if node is None:
            self._size += 1
            return self.Node(key, val)
        if node.key > key:
            node.left = self._add(node.left, key, val)
        elif node.key < key:
            node.right = self._add(node.right, key, val)
        else:
            # key相等，更新val的值
            node.val = val

        # 维护红黑树的性质
        if self._is_red(node.right) and not self._is_red(node.right):
            node = self._left_rotate(node)

        if self._is_red(node.left) and self._is_red(node.left.left):
            node = self._right_rotate(node)

        if self._is_red(node.left) and self._is_red(node.right):
            self._flip_colors(node)

        return node

    # def remove(self, key):
    #     """
    #     @删除AVLTree中键为key的节点
    #     @返回删除节点的val(没有时返回None)
    #     """
    #     node = self.get_node(self.root, key)
    #     if node is not None:
    #         self.root = self._remove(self.root, key)
    #         return node.val
    #     return None
    #
    # def _remove(self, node, key):
    #     """
    #     @删除AVLTree中键为key的节点， 递归算法
    #     @返回删除节点后的新的AVLTree的根节点
    #     """
    #     if node is None:
    #         return None
    #     if node.key > key:
    #         node.left = self._remove(node.left, key)
    #         ret = node
    #     elif node.key < key:
    #         node.right = self._remove(node.right, key)
    #         ret = node
    #     else:  # node.key == key
    #         if node.left is None:
    #             right_node = node.right
    #             node.right = None
    #             self._size -= 1
    #             ret = right_node
    #         elif node.right is None:
    #             left_node = node.left
    #             node.left = None
    #             self._size -= 1
    #             ret = left_node
    #         else:
    #             # 左右子树均不为空的情况， 用右子树的最小值替代要删除的那个节点，也可以用左子树的最大值替代
    #             successor = self._minimum(node.right)
    #             # 这里原来使用_remove_min来删除最小节点，但是可能导致平衡性破坏，故使用_remove代替
    #             successor.right = self._remove(node.right, successor.key)
    #             # self._size += 1  # 注意，这里由于删除的最小节点移动到前面去了，相当于没有删除，需要将_remove中的size减一操作撤销
    #             # 这里_remove(node.right, successor.key)肯定会将size减一，因为右子树的最小值节点坑定没有左子树，所以最终一定是通过第466
    #             # 行的if分支进行节点删除操作，在那里进行_size减一操作
    #             successor.left = node.left
    #             node.left = None   # 移除node，这里真正需要将size减一
    #             node.right = None
    #             # 前面size加一，后面size减一，所以抵消，全部注释掉，实际上相当一size减一，只不过是在_remove中完成的
    #             # self._size -= 1
    #             ret = successor
    #
    #     # 删除节点后得到的ret可能为空
    #     if ret is None:
    #         return ret
    #
    #     # 更新height
    #     # 插入新节点可能导致高度值发生变化，对当前节点的高度值进行重新计算
    #     ret.height = 1 + max(self._get_height(ret.left), self._get_height(ret.right))
    #
    #     # 计算节点的平衡因子
    #     balance_factor = self._get_balance_factor(ret)
    #
    #     # 维护平衡性
    #     # LL
    #     if balance_factor > 1 and self._get_balance_factor(ret.left) >= 0:
    #         return self.right_rotate(ret)
    #
    #     # RR
    #     if balance_factor < -1 and self._get_balance_factor(ret.right) <= 0:
    #         return self.left_rotate(ret)
    #
    #     # LR
    #     if balance_factor > 1 and self._get_balance_factor(ret.left) < 0:
    #         ret.left = self.left_rotate(ret.left)
    #         return self.right_rotate(ret)
    #
    #     # RL
    #     if balance_factor < -1 and self._get_balance_factor(ret.right) > 0:
    #         ret.right = self.right_rotate(ret.right)
    #         return self.left_rotate(ret)
    #
    #     return ret